package com.xzzz.common.lb;

import com.xzzz.common.base.pojo.R;
import com.xzzz.common.lb.pojo.ServicesLbRes;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.discovery.DiscoveryClient;
import org.springframework.cloud.client.loadbalancer.reactive.ReactiveLoadBalancer;
import org.springframework.cloud.loadbalancer.support.LoadBalancerClientFactory;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.ArrayList;
import java.util.List;

/**
 * 负载均衡
 *
 * @author xzzz
 */
@RestController
@RequestMapping("/lb")
public class LoadBalancerController {

    @Autowired
    private LoadBalancerClientFactory loadBalancerClientFactory;

    @Autowired(required = false)
    private DiscoveryClient discoveryClient;

    /**
     * 获取LB策略
     *
     * @param serviceId 服务ID
     * @return 负载均衡策略
     */
    @GetMapping
    public R<String> getLoadBalancer(@RequestParam String serviceId) {
        ReactiveLoadBalancer<ServiceInstance> loadBalancer = loadBalancerClientFactory.getInstance(serviceId);
        return R.ok(loadBalancer.getClass().getSimpleName());
    }

    /**
     * 获取所有服务的所有LB策略
     */
    @GetMapping("/all")
    public R<List<ServicesLbRes>> getLoadBalancerAll() {
        List<ServicesLbRes> loadBalancerAll = new ArrayList<>();

        List<String> serviceIds = discoveryClient.getServices();
        for (String serviceId : serviceIds) {
            ReactiveLoadBalancer<ServiceInstance> loadBalancer = loadBalancerClientFactory.getInstance(serviceId);
            if (loadBalancer != null) {
                ServicesLbRes res = new ServicesLbRes();
                res.setServiceId(serviceId);
                res.setLoadBalancer(loadBalancer.getClass().getSimpleName());
                loadBalancerAll.add(res);
            }
        }
        return R.ok(loadBalancerAll);
    }

    public R<?> resetLB(String serviceId) {
        return R.ok();
    }
}
